import time
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import unittest

class TestLogin(unittest.TestCase):
    def setUp(self):
        self.driver = webdriver.Edge()
        self.driver.implicitly_wait(10)
        self.base_url = 'http://127.0.0.1:5000'
        
    def tearDown(self):
        if hasattr(self, 'driver'):
            self.driver.quit()

    def _login_user(self, username_or_email, password):
        """登录用户的辅助方法"""
        self.driver.get(f'{self.base_url}/login')
        time.sleep(1)
        
        username_input = self.driver.find_element(By.ID, 'username_or_email')
        username_input.send_keys(username_or_email)
        
        password_input = self.driver.find_element(By.ID, 'password')
        password_input.send_keys(password)
        
        submit_button = self.driver.find_element(By.CLASS_NAME, 'btn')
        submit_button.click()
        time.sleep(1)

    def test_login_success(self):
        """测试登录成功的情况"""
        self._login_user('testuser123', '123456')
        
        # 验证是否跳转到首页
        try:
            WebDriverWait(self.driver, 5).until(
                EC.url_to_be(f'{self.base_url}/')
            )
            self.assertEqual(self.driver.current_url, f'{self.base_url}/')
        except TimeoutException:
            self.fail("登录失败，未跳转到首页")

    def test_login_failure_wrong_password(self):
        """测试登录失败 - 密码错误"""
        self._login_user('testuser123', 'wrongpassword')
        
        # 验证是否停留在登录页面
        self.assertEqual(self.driver.current_url, f'{self.base_url}/login')
        time.sleep(2)

    def test_login_failure_nonexistent_user(self):
        """测试登录失败 - 用户不存在"""
        self._login_user('nonexistentuser', '123456')
        
        # 验证是否停留在登录页面
        self.assertEqual(self.driver.current_url, f'{self.base_url}/login')

if __name__ == '__main__':
    unittest.main() 